
from llama_index.core import (
    VectorStoreIndex,
    StorageContext,
    load_index_from_storage,
    ServiceContext
)
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.elasticsearch import ElasticsearchStore
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import (
    ResponseMode,
    get_response_synthesizer,
)
from llama_index.llms.openai import OpenAI
from fastapi import APIRouter
import yaml

router = APIRouter(
    prefix="/gj-rag",
    tags=["query"],
    responses={404: {"description": "Not found"}}
)

# 配置
with open('config.yaml', 'r') as file:
        config = yaml.load(file, Loader=yaml.FullLoader)

es = ElasticsearchStore(
    index_name=config["question_es"]["index_name"],
    es_url=config["question_es"]["url"],
    es_user=config["question_es"]["user"],
    es_password=config["question_es"]["password"]
)

llm = OpenAI(model=config["openai"]["model"], api_base=config["openai"]["api_base"], api_key=config["openai"]["api_key"])
embed_model = OpenAIEmbedding(api_base=config["openai"]["api_base"], api_key=config["openai"]["api_key"])

service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
index = VectorStoreIndex.from_vector_store(vector_store=es, service_context=service_context)


retriever = index.as_retriever(
    similarity_top_k=5,
)


def querstion_query(query:str):
    response = retriever.retrieve(query)
    result = []
    for node_with_score in response:
        # 获取节点的text属性
        text_content = node_with_score.node.text
        result.append(text_content)
    return result


@router.get("/query/{query_string}")
async def query_endpoint(query_string: str):
    return {"result": querstion_query(query_string)}

# if __name__ == "__main__":
#     import uvicorn
# #     import asyncio
# #     loop = asyncio.get_event_loop()
#     uvicorn.run(app, host="0.0.0.0", port=8003)


# 启动代码
    # python3 -m uvicorn question_query:app --host 0.0.0.0 --port 8003  --loop asyncio